import torch
from transformers import SpeechT5Processor, SpeechT5ForTextToSpeech, SpeechT5HifiGan
from IPython.display import Audio

# -------------------------------
# 1. 定义本地检查点路径
# -------------------------------
checkpoint_path = "checkpoint-100"  # 你的模型保存目录

# -------------------------------
# 2. 加载 Processor 和 Model
# -------------------------------
print(f"正在从 {checkpoint_path} 加载模型...")

# 注意：必须先加载 processor，再加载 model
processor = SpeechT5Processor.from_pretrained(checkpoint_path)
model = SpeechT5ForTextToSpeech.from_pretrained(checkpoint_path)

# 将模型移到 GPU（如果可用）
device = "cuda" if torch.cuda.is_available() else "cpu"
model = model.to(device)

# -------------------------------
# 3. 加载声码器（Vocoder）—— 用于将频谱图转为音频波形
# -------------------------------
# 注意：vocoder 不在 checkpoint 中，需单独加载
vocoder = SpeechT5HifiGan.from_pretrained("microsoft/speecht5_hifigan").to(device)

# -------------------------------
# 4. 准备输入文本
# -------------------------------
text = "hallo allemaal, ik praat nederlands. groetjes aan iedereen!"

# -------------------------------
# 5. 获取说话人嵌入（Speaker Embedding）
# -------------------------------
# 假设你从测试集中取一个样本
example = final_dataset["test"][304]
speaker_embeddings = torch.tensor(example["speaker_embeddings"]).unsqueeze(0)  # -> (1, 512)
speaker_embeddings = speaker_embeddings.to(device)

# -------------------------------
# 6. 文本编码
# -------------------------------
inputs = processor(text=text, return_tensors="pt").to(device)
# inputs 包含: input_ids, attention_mask

# -------------------------------
# 7. 推理生成频谱图
# -------------------------------
with torch.no_grad():
    spectrogram = model.generate_speech(
        inputs["input_ids"],
        speaker_embeddings=speaker_embeddings,
        attention_mask=inputs["attention_mask"]
    )

# -------------------------------
# 8. 使用声码器生成音频波形
# -------------------------------
with torch.no_grad():
    waveform = vocoder(spectrogram)  # (1, T)

# -------------------------------
# 9. 转换为 numpy 并播放
# -------------------------------
waveform = waveform.cpu().numpy().squeeze()  # 去掉 batch 维度
sampling_rate = 16000  # SpeechT5 默认采样率

# 播放音频
Audio(waveform, rate=sampling_rate)